#pragma once

#include "onnxruntime_cxx_api.h"

#include "ep_utils.h"
#include "flatbuffers/idl.h"
#include "ort_trt_int8_cal_table.fbs.h"
#include "make_string.h"

#include "nv_includes.h"
#include "gsl/narrow"

#include <fstream>
#include <unordered_map>
#include <string>
#include <vector>
#include <sstream>
#include <iostream>
#include <filesystem>

namespace fs = std::filesystem;

namespace trt_ep {

bool CalcMemSizeForArrayWithAlignment(size_t nmemb, size_t size, size_t alignment, size_t* out) noexcept {
  size_t alloc_size = size;
  if (alignment == 0) {
    *out = alloc_size * nmemb;
  } else {
    size_t alignment_mask = alignment - 1;
    *out = (alloc_size * nmemb + alignment_mask) & ~static_cast<size_t>(alignment_mask);
  }
  return true;
}

template <typename T>
AllocatorUniquePtr<T> MakeUniquePtrFromOrtAllocator(OrtAllocator* ort_allocator, size_t count_or_bytes,
                                                    bool use_reserve = false) {
  size_t alloc_size = count_or_bytes;
  // if T is not void, 'count_or_bytes' == number of items so allow for that
  if constexpr (!std::is_void<T>::value) {
    // sizeof(void) isn't valid, but the compiler isn't smart enough to ignore that this line isn't
    // reachable if T is void. use std::conditional to 'use' void* in the sizeof call
    constexpr auto size = sizeof(typename std::conditional<std::is_void<T>::value, void*, T>::type);
    CalcMemSizeForArrayWithAlignment(count_or_bytes, size, 0, &alloc_size);
  }

  T* p = nullptr;
  if (use_reserve) {
    p = static_cast<T*>(ort_allocator->Reserve(ort_allocator, alloc_size));
  } else {
    p = static_cast<T*>(ort_allocator->Alloc(ort_allocator, alloc_size));
  }

  return AllocatorUniquePtr<T>{p, [ort_allocator](T* p) { ort_allocator->Free(ort_allocator, p); }};
}

// Following helper functions/struct, GetNodeInputEdgeCount, GetOutputNodes, KahnsTopologicalSort, VisitorPriorityQueue, PriorityNodeCompare are added but are not used for now.
// TODO: They will be used for graph partition in the following PR.

template <typename T>
struct VisitorPriorityQueue {
  using ComparatorType = std::function<bool(T, T)>;
  std::list<T> list_;
  const ComparatorType comparator_ = nullptr;
  VisitorPriorityQueue(const ComparatorType& comp) : comparator_(comp) {}

  void push(T node) {
    list_.insert(
        std::upper_bound(list_.begin(), list_.end(), node, comparator_),
        node);
  }
  bool empty() { return list_.empty(); }
  T top() { return list_.back(); }
  void pop() { list_.pop_back(); }
};

// Get the number of input edges that come from another node upstream.
static OrtStatus* GetNodeInputEdgeCount(const OrtNode* node, size_t& num_input_edges) {
  const OrtApi& ort_api = Ort::GetApi();

  size_t num_inputs = 0;
  RETURN_IF_ERROR(ort_api.Node_GetNumInputs(node, &num_inputs));

  std::vector<const OrtValueInfo*> inputs(num_inputs);
  RETURN_IF_ERROR(ort_api.Node_GetInputs(node, inputs.data(), inputs.size()));

  // Sum the number of inputs with a producer node.
  num_input_edges = 0;

  for (const OrtValueInfo* ort_input : inputs) {
    Ort::ConstValueInfo input{ort_input};
    if (input == nullptr) continue;  // Skip missing optional input

    auto producer_info = input.GetProducerNode();
    num_input_edges += static_cast<size_t>(producer_info.node != nullptr);
  }

  return nullptr;
}

// Get all output nodes that consume an output from the given node.
static OrtStatus* GetOutputNodes(const OrtNode* node, std::vector<Ort::ConstNode>& result) {
  const OrtApi& ort_api = Ort::GetApi();

  size_t num_outputs = 0;
  RETURN_IF_ERROR(ort_api.Node_GetNumOutputs(node, &num_outputs));

  std::vector<const OrtValueInfo*> outputs(num_outputs);
  RETURN_IF_ERROR(ort_api.Node_GetOutputs(node, outputs.data(), outputs.size()));

  std::vector<Ort::ConstNode> output_nodes;
  output_nodes.reserve(num_outputs);  // May have more than `num_outputs`

  // Gather the OrtNode consumers of every output.
  for (const OrtValueInfo* ort_output : outputs) {
    Ort::ConstValueInfo output{ort_output};
    if (output == nullptr) continue;  // Skip missing optional output

    auto consumers_info = output.GetConsumers();
    for (const auto& consumer : consumers_info) {
      output_nodes.push_back(consumer.node);
    }
  }

  result = std::move(output_nodes);
  return nullptr;
}

// Kahn's topological sort.
// Adapted from onnxruntime/core/graph/graph.cc to use public C API graph types.
static OrtStatus* KahnsTopologicalSort(const OrtGraph& graph,
                                       const std::function<void(const OrtNode*)>& enter,
                                       const std::function<bool(const OrtNode*, const OrtNode*)>& comp) {
  const OrtApi& ort_api = Ort::GetApi();

  try {
    // Get all nodes
    size_t num_nodes = 0;
    RETURN_IF_ERROR(ort_api.Graph_GetNumNodes(&graph, &num_nodes));

    if (num_nodes == 0) {
      return Ort::Status{nullptr};  // Nothing to sort.
    }

    std::vector<const OrtNode*> nodes(num_nodes);
    RETURN_IF_ERROR(ort_api.Graph_GetNodes(&graph, nodes.data(), nodes.size()));

    // Get the maximum node ID. Not really required if we chose to represent the `in_degree` as a map instead of vector.
    size_t max_node_id = 0;
    for (const OrtNode* node : nodes) {
      size_t node_id = 0;
      RETURN_IF_ERROR(ort_api.Node_GetId(node, &node_id));
      max_node_id = std::max(max_node_id, node_id);
    }

    std::vector<size_t> in_degree(max_node_id + 1, 0);
    std::vector<size_t> topo_order;
    VisitorPriorityQueue<const OrtNode*> to_visit(comp);

    topo_order.reserve(num_nodes);

    // Initialize in_degree and initial nodes to visit first.
    for (const OrtNode* node : nodes) {
      size_t input_edge_count = 0;
      RETURN_IF_ERROR(GetNodeInputEdgeCount(node, input_edge_count));

      size_t node_id = 0;
      RETURN_IF_ERROR(ort_api.Node_GetId(node, &node_id));

      in_degree[node_id] = input_edge_count;
      if (input_edge_count == 0) {
        to_visit.push(node);
      }
    }

    while (!to_visit.empty()) {
      const OrtNode* current_node = to_visit.top();
      to_visit.pop();

      if (!current_node) continue;

      if (enter) {
        enter(current_node);
      }

      std::vector<Ort::ConstNode> output_nodes;
      RETURN_IF_ERROR(GetOutputNodes(current_node, output_nodes));

      for (const auto& output_node : output_nodes) {
        size_t output_node_id = 0;
        RETURN_IF_ERROR(ort_api.Node_GetId(output_node, &output_node_id));

        auto& node_in_degree = in_degree[output_node_id];
        node_in_degree--;

        if (node_in_degree == 0) {
          to_visit.push(output_node);
        }
      }

      size_t current_node_id = 0;
      RETURN_IF_ERROR(ort_api.Node_GetId(current_node, &current_node_id));
      topo_order.push_back(current_node_id);
    }

    if (num_nodes != topo_order.size()) {
      return Ort::Status("Some nodes are not included in the topological sort: graph has a cycle", ORT_FAIL);
    }
  } catch (const Ort::Exception& ex) {
    Ort::Status status(ex);
    return status.release();
  } catch (const std::exception& ex) {
    Ort::Status status(ex.what(), ORT_EP_FAIL);
    return status.release();
  }

  return nullptr;
}

// Node comparison functor copied from onnxruntime/core/graph/graph.cc
struct PriorityNodeCompare {
  inline bool IsHighPri(const OrtNode* n) const {
    // local statics so we can compare std::strings in the checks
    static constexpr std::string_view shape_op("Shape");
    static constexpr std::string_view size_op("Size");

    const char* op_type = nullptr;
    Ort::Status status(Ort::GetApi().Node_GetOperatorType(n, &op_type));
    ENFORCE(status.IsOK());

    return shape_op == op_type || size_op == op_type;
  }

  // Used for std::priority_queue
  // If return false, n1 will be output first
  // If return true, n2 will be output first
  bool operator()(const OrtNode* n1, const OrtNode* n2) const {
    // nodes in global high priority list will be output first
    const bool isN1HighPri = IsHighPri(n1);
    const bool isN2HighPri = IsHighPri(n2);
    if (isN1HighPri != isN2HighPri) {
      return isN2HighPri;
    }

    // nodes with lower priority value will be output first
    const auto n1_priority = 0;  // n1->Priority(); // Looks to always be 0 inside ORT?
    const auto n2_priority = 0;  // n2->Priority(); // Looks to always be 0 inside ORT?
    if (n1_priority != n2_priority) {
      return n1_priority > n2_priority;
    }

    // otherwise, nodes with lower index will be output first
    size_t n1_id = 0;
    Ort::Status status1(Ort::GetApi().Node_GetId(n1, &n1_id));
    ENFORCE(status1.IsOK());

    size_t n2_id = 0;
    Ort::Status status2(Ort::GetApi().Node_GetId(n2, &n2_id));
    ENFORCE(status2.IsOK());

    return n1_id > n2_id;
  }
};

bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<std::string, float>& dynamic_range_map) {
  // Set dynamic range for input tensors
  for (int i = 0; i < network.getNbInputs(); ++i) {
    const std::string tensor_name = network.getInput(i)->getName();
    auto dynamic_range_iter = dynamic_range_map.find(tensor_name);
    if (dynamic_range_iter != dynamic_range_map.end()) {
#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
      if (!network.getInput(i)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) {
#if defined(_MSC_VER)
#pragma warning(pop)
#endif
        // LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for network input " << tensor_name;
        return false;
      }
    }
  }

  // Set dynamic range for activations and weights
  for (int i = 0; i < network.getNbLayers(); ++i) {
    auto trt_layer = network.getLayer(i);
    for (int j = 0, e = trt_layer->getNbOutputs(); j < e; ++j) {
      const std::string tensor_name = trt_layer->getOutput(j)->getName();
      auto dynamic_range_iter = dynamic_range_map.find(tensor_name);
      if (dynamic_range_iter != dynamic_range_map.end()) {
#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
        if (!trt_layer->getOutput(j)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) {
#if defined(_MSC_VER)
#pragma warning(pop)
#endif
          // LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for tensor " << tensor_name;
          return false;
        }
      } else if (trt_layer->getType() == nvinfer1::LayerType::kCONSTANT) {
        nvinfer1::IConstantLayer* const_layer = static_cast<nvinfer1::IConstantLayer*>(trt_layer);
        const std::string const_layer_name = const_layer->getName();
        auto trt_weights = const_layer->getWeights();
        double max_weight = std::numeric_limits<double>::min();
        for (int64_t k = 0, end = trt_weights.count; k < end; ++k) {
          double weight{};
          switch (trt_weights.type) {
            case nvinfer1::DataType::kFLOAT:
              weight = static_cast<const float*>(trt_weights.values)[k];
              break;
            case nvinfer1::DataType::kBOOL:
              weight = static_cast<const bool*>(trt_weights.values)[k];
              break;
            case nvinfer1::DataType::kINT8:
              weight = static_cast<const int8_t*>(trt_weights.values)[k];
              break;
            case nvinfer1::DataType::kHALF:
              weight = static_cast<const uint16_t*>(trt_weights.values)[k];
              break;
            case nvinfer1::DataType::kINT32:
              weight = static_cast<const int32_t*>(trt_weights.values)[k];
              break;
#if NV_TENSORRT_MAJOR >= 10
            case nvinfer1::DataType::kINT64:
              weight = static_cast<double>(static_cast<const int64_t*>(trt_weights.values)[k]);
              break;
#endif  // NV_TENSORRT_MAJOR >= 10
            default:
              // LOGS_DEFAULT(ERROR) << "Found unsupported datatype for layer " << const_layer_name;
              return false;
          }
          max_weight = std::max(max_weight, std::abs(weight));
        }
#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4996)
#endif
        if (!trt_layer->getOutput(j)->setDynamicRange(static_cast<float>(-max_weight),
                                                      static_cast<float>(max_weight))) {
#if defined(_MSC_VER)
#pragma warning(pop)
#endif
          // LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for layer " << const_layer_name;
          return false;
        }
      }
    }
  }
  return true;
}

std::vector<std::string> SplitToStringVec(std::string const& s, char separator) {
  std::vector<std::string> splitted;

  for (size_t start = 0; start < s.length();) {
    size_t separatorIndex = s.find(separator, start);
    if (separatorIndex == std::string::npos) {
      separatorIndex = s.length();
    }
    splitted.emplace_back(s.substr(start, separatorIndex - start));
    start = separatorIndex + 1;
  }

  return splitted;
}

nvinfer1::TacticSources GetTacticSourceFromString(std::string& tactic_string) {
  nvinfer1::TacticSources disabledTactics = 0;
  nvinfer1::TacticSources enabledTactics = 0;
  std::vector<std::string> tacticList = SplitToStringVec(tactic_string, ',');
  for (auto& t : tacticList) {
    bool enable{false};
    if (t.front() == '+') {
      enable = true;
    } else if (t.front() != '-') {
      // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic source must be prefixed with + or - skipping: " << t;
    }
    t.erase(0, 1);

    const auto toUpper = [](std::string& sourceName) {
      std::transform(sourceName.begin(), sourceName.end(), sourceName.begin(),
                     [](char c) { return gsl::narrow<char>(std::toupper(c)); });
      return sourceName;
    };

    nvinfer1::TacticSource source{};
    t = toUpper(t);
    if (t == "CUBLAS") {
      // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic kCUBLAS is deprecated in TensorRT 10.0";
#if NV_TENSORRT_MAJOR < 10
      source = nvinfer1::TacticSource::kCUBLAS;
#endif
    } else if (t == "CUBLASLT" || t == "CUBLAS_LT") {
      // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic kCUBLAS_LT is deprecated in TensorRT 9.0";
#if NV_TENSORRT_MAJOR < 9
      source = nvinfer1::TacticSource::kCUBLAS_LT;
#endif
    } else if (t == "CUDNN") {
      // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic kCUDNN is deprecated in TensorRT 10.0";
#if NV_TENSORRT_MAJOR < 10
      source = nvinfer1::TacticSource::kCUDNN;
#endif
    } else if (t == "EDGE_MASK_CONVOLUTIONS") {
      source = nvinfer1::TacticSource::kEDGE_MASK_CONVOLUTIONS;
    } else if (t == "JIT_CONVOLUTIONS") {
      source = nvinfer1::TacticSource::kJIT_CONVOLUTIONS;
    } else {
      // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Tactic source was not found with name: " << t;
    }

    uint32_t sourceBit = 1U << static_cast<uint32_t>(source);

    if (enable) {
      enabledTactics |= sourceBit;
    } else {
      disabledTactics |= sourceBit;
    }
  }
  return enabledTactics & ~disabledTactics;
}

inline std::vector<char> loadTimingCacheFile(const std::string inFileName) {
  std::ifstream iFile(inFileName, std::ios::in | std::ios::binary);
  if (!iFile) {
    // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Could not read timing cache from: " << inFileName
    //                       << ". A new timing cache will be generated and written.";
    return std::vector<char>();
  }
  iFile.seekg(0, std::ifstream::end);
  size_t fsize = iFile.tellg();
  iFile.seekg(0, std::ifstream::beg);
  std::vector<char> content(fsize);
  iFile.read(content.data(), fsize);
  iFile.close();
  return content;
}

inline void saveTimingCacheFile(const std::string outFileName, const nvinfer1::IHostMemory* blob) {
  std::ofstream oFile(outFileName, std::ios::out | std::ios::binary);
  if (!oFile) {
    // LOGS_DEFAULT(WARNING) << "[TensorRT EP] Could not write timing cache to: " << outFileName;
    return;
  }
  oFile.write((char*)blob->data(), blob->size());
  oFile.close();
}

float ConvertSinglePrecisionIEEE754ToFloat(unsigned long input) {
  int s = (input >> 31) & 0x01;
  int e = ((input & 0x7f800000) >> 23) - 127;
  int p = -1;
  double m = 0.0;
  for (int i = 0; i < 23; ++i) {
    m += ((input >> (23 - i - 1)) & 0x01) * pow(2.0, p--);
  }
  return static_cast<float>((s ? -1 : 1) * pow(2.0, e) * (m + 1.0));
}

/*
 * Read calibration table for INT8 quantization
 * Two kind of calibration tables are supported,
 * 1. ORT generated calibration table
 * The table is pre-serialized by flatbuffers.
 * Each entry in the table is a key-value pair,
 * key: tensor name, value: maximum absolute value in floating point
 * For example,
 *   data_0 2.008338
 *   ...
 * 2. Native TensorRT generated calibration table
 * Data format is defined by TensorRT as,
 * tensor name : scale in 32-bit single precision IEEE754 format
 * For example,
 *   TRT-7103-EntropyCalibration2
 *   data_0: 4000889d
 *   ...
 */
bool ReadDynamicRange(const std::string file_name, const bool is_trt_calibration_table, std::unordered_map<std::string, float>& dynamic_range_map) {
  std::ifstream infile(file_name, std::ios::binary | std::ios::in);
  if (!infile) {
    return false;
  }

  if (is_trt_calibration_table) {
    // Native TensorRT generated calibration table
    std::string line;
    char delim = ':';
    if (std::getline(infile, line)) {
      std::istringstream first_line(line);
      std::string version;
      std::getline(first_line, version, delim);
      std::size_t found = version.find("TRT-");
      if (found != std::string::npos) {
        while (std::getline(infile, line)) {
          std::istringstream in_line(line);
          std::string str;
          std::getline(in_line, str, delim);
          std::string tensor_name = str;
          std::getline(in_line, str, delim);
          unsigned long scale_int = std::strtoul(str.c_str(), nullptr, 16);
          float scale_float = ConvertSinglePrecisionIEEE754ToFloat(scale_int);
          float dynamic_range = scale_float * 127.0f;
          dynamic_range_map[tensor_name] = dynamic_range;
        }
      } else {
        throw std::runtime_error("This is not a TensorRT generated calibration table " + file_name);
      }
    }
  } else {
    // ORT generated calibration table
    infile.seekg(0, std::ios::end);
    size_t length = infile.tellg();
    infile.seekg(0, std::ios::beg);
    std::unique_ptr<char[]> data{new char[length]};
    infile.read((char*)data.get(), length);
    infile.close();
    auto flat_table = flatbuffers::GetRoot<CalTableFlatBuffers::TrtTable>((const uint8_t*)data.get());
    auto flat_dict = flat_table->dict();
    for (size_t i = 0, end = flat_dict->size(); i < end; ++i) {
      flatbuffers::uoffset_t idx = static_cast<flatbuffers::uoffset_t>(i);
      dynamic_range_map[flat_dict->Get(idx)->key()->str()] = std::stof(flat_dict->Get(idx)->value()->str());
    }
  }
  return true;
}

/*
 * Get number of profile setting.
 *
 * profile_min_shapes/profile_max_shapes/profile_opt_shapes may contain multiple profile settings.
 * Note: TRT EP currently only supports one profile setting.
 *
 * {
 *   tensor_a: [[dim_0_value_0, dim_1_value_1, dim_2_value_2]],
 *   tensor_b: [[dim_0_value_3, dim_1_value_4, dim_2_value_5]]
 * }
 *
 */
int GetNumProfiles(std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_shapes) {
  int num_profile = 0;
  for (auto it = profile_shapes.begin(); it != profile_shapes.end(); it++) {
    num_profile = static_cast<int>(it->second.size());
    if (num_profile > 0) {
      break;
    }
  }
  return num_profile;
}

/*
 * Seralize engine profile
 * The profile contains min/max shape ranges of dynamic shape dimensions of each input tensor
 * For example, assume tensor_a has two dynamic shape dimensions: dim_0 and dim_2, and tensor_b
 * has one dynamic shape dimension: dim_1. The data in profile will be,
 * key: tensor_a, value: dim_0 min_shape max_shape dim_2 min_shape max_shape
 * key: tensor_b, value: dim_1 min_shape max_shape
 *
 * [Deprecated] Use SerializeProfileV2
 */
void SerializeProfile(const std::string& file_name, std::unordered_map<std::string, std::unordered_map<size_t, std::pair<int64_t, int64_t>>>& shape_ranges) {
  // Serialize profile
  flexbuffers::Builder builder;
  auto profile_start = builder.StartMap();
  for (auto outer_it = shape_ranges.begin(); outer_it != shape_ranges.end(); ++outer_it) {
    builder.TypedVector(outer_it->first.c_str(), [&] {
      for (auto inner_it = outer_it->second.begin(); inner_it != outer_it->second.end(); ++inner_it) {
        builder.Int(inner_it->first);
        builder.Int(inner_it->second.first);
        builder.Int(inner_it->second.second);
      }
    });
  }
  builder.EndMap(profile_start);
  builder.Finish();

  // Save flexbuffer
  std::ofstream file(file_name, std::ios::binary | std::ios::out);
  auto buf = builder.GetBuffer();
  size_t size = builder.GetSize();
  file.write(reinterpret_cast<const char*>(&buf[0]), size);
  file.close();
}

// Deserialize engine profile
// [Deprecated] Use DeserializeProfileV2
std::unordered_map<std::string, std::unordered_map<size_t, std::pair<int64_t, int64_t>>> DeserializeProfile(std::ifstream& infile) {
  // Load flexbuffer
  infile.seekg(0, std::ios::end);
  size_t length = infile.tellg();
  infile.seekg(0, std::ios::beg);
  std::unique_ptr<char[]> data{new char[length]};
  infile.read((char*)data.get(), length);
  infile.close();

  // Deserialize profile
  std::unordered_map<std::string, std::unordered_map<size_t, std::pair<int64_t, int64_t>>> shape_ranges;
  auto tensors_range_entries = flexbuffers::GetRoot((const uint8_t*)data.get(), length).AsMap();
  auto keys = tensors_range_entries.Keys();
  auto values = tensors_range_entries.Values();
  for (size_t i = 0, i_end = keys.size(); i < i_end; ++i) {
    auto dim_range_vectors = values[i].AsTypedVector();
    std::unordered_map<size_t, std::pair<int64_t, int64_t>> inner_map;
    for (size_t j = 0, j_end = dim_range_vectors.size() / 3; j < j_end; ++j) {
      size_t idx = 3 * j;
      inner_map[dim_range_vectors[idx].AsInt64()] = std::make_pair(dim_range_vectors[idx + 1].AsInt64(), dim_range_vectors[idx + 2].AsInt64());
    }
    shape_ranges[keys[i].AsString().c_str()] = inner_map;
  }
  return shape_ranges;
}

/*
 * Seralize engine profile. (This function starts from ORT 1.15)
 *
 *
 * (1) Single profile case:
 * Assume tensor_a has two dynamic shape dimensions: dim_0 and dim_2,
 * and tensor_b has one dynamic shape dimension: dim_1.
 *
 * The data before serialization will be:
 * {
 *   tensor_a: {
 *     dim_0: [[min_shape_0, max_shape_0, opt_shape_0]],
 *     dim_2: [[min_shape_2, max_shape_2, opt_shape_2]]
 *   },
 *   tensor_b: {
 *     dim_1: [[min_shape_1, max_shape_1, opt_shape_1]]
 *   }
 * }
 *
 * The data after serialization will be:
 * {
 *   tensor_a: [dim_0, min_shape_0, max_shape_0, opt_shape_0, dim_2, min_shape_2, max_shape_2, opt_shape_2]
 *   tensor_b: [dim_1, min_shape_1, max_shape_1, opt_shape_1]
 * }
 *
 *
 * (2) Multiple profiles case:
 * For example, if the data before serialization is:
 * {
 *   tensor_a: {
 *     dim_0: [[min_shape_0, max_shape_0, opt_shape_0], [min_shape_1, max_shape_1, opt_shape_1]]
 *   },
 *   tensor_b: {
 *     dim_1: [[min_shape_2, max_shape_2, opt_shape_2], [min_shape_3, max_shape_3, opt_shape_3]]
 *   }
 * }
 *
 * The data after serialization will be:
 * {
 *   tensor_a: [dim_0, min_shape_0, max_shape_0, opt_shape_0, dim_0, min_shape_1, max_shape_1, opt_shape_1]
 *              |                                          |  |                                          |
 *              ---------------- profile 0 -----------------  ---------------- profile 1 -----------------
 *
 *   tensor_b: [dim_1, min_shape_2, max_shape_2, opt_shape_2, dim_1, min_shape_3, max_shape_3, opt_shape_3]
 *              |                                          |  |                                          |
 *              ---------------- profile 0 -----------------  ---------------- profile 1 -----------------
 * }
 *
 */
void SerializeProfileV2(const std::string& file_name, std::unordered_map<std::string, std::unordered_map<size_t, std::vector<std::vector<int64_t>>>>& shape_ranges) {
  ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] In SerializeProfileV2()";
  // Serialize profile
  flexbuffers::Builder builder;
  auto tensor_map_start = builder.StartMap();
  for (auto tensor_it = shape_ranges.begin(); tensor_it != shape_ranges.end(); tensor_it++) {  // iterate tensors
    ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] input tensor is '" << tensor_it->first.c_str() << "'";
    builder.TypedVector(tensor_it->first.c_str(), [&] {
      for (auto dim_it = tensor_it->second.begin(); dim_it != tensor_it->second.end(); dim_it++) {
        size_t num_profiles = dim_it->second.size();
        for (size_t i = 0; i < num_profiles; i++) {
          ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] profile #" << i << ", dim is " << dim_it->first;
          builder.Int(dim_it->first);
          builder.Int(dim_it->second[i][0]);
          builder.Int(dim_it->second[i][1]);
          builder.Int(dim_it->second[i][2]);
          ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] " << dim_it->first << ", " << dim_it->second[i][0] << ", " << dim_it->second[i][1] << ", " << dim_it->second[i][2];
        }
      }
    });
  }
  builder.EndMap(tensor_map_start);
  builder.Finish();

  // Save flexbuffer
  std::ofstream file(file_name, std::ios::binary | std::ios::out);
  auto buf = builder.GetBuffer();
  size_t size = builder.GetSize();
  file.write(reinterpret_cast<const char*>(&buf[0]), size);
  file.close();
}

/*
 * Deserialize engine profile. (This function starts from ORT 1.15)
 *
 *
 * (1) Single profile case:
 * Assume tensor_a has two dynamic shape dimensions: dim_0 and dim_2,
 * and tensor_b has one dynamic shape dimension: dim_1.
 *
 * The data in profile file will be:
 * {
 *   tensor_a: [dim_0, min_shape_0, max_shape_0, opt_shape_0, dim_2, min_shape_2, max_shape_2, opt_shape_2]
 *   tensor_b: [dim_1, min_shape_1, max_shape_1, opt_shape_1]
 * }
 *
 * The data after deserialization will be:
 * {
 *   tensor_a: {
 *     dim_0: [[min_shape_0, max_shape_0, opt_shape_0]],
 *     dim_2: [[min_shape_2, max_shape_2, opt_shape_2]]
 *   },
 *   tensor_b: {
 *     dim_1: [[min_shape_1, max_shape_1, opt_shape_1]]
 *   }
 * }
 *
 *
 * (2) Multiple profiles case:
 * For example, if the data in profile file is:
 * {
 *   tensor_a: [dim_0, min_shape_0, max_shape_0, opt_shape_0, dim_0, min_shape_1, max_shape_1, opt_shape_1]
 *              |                                          |  |                                          |
 *              ---------------- profile 0 -----------------  ---------------- profile 1 -----------------
 *
 *   tensor_b: [dim_1, min_shape_2, max_shape_2, opt_shape_2, dim_1, min_shape_3, max_shape_3, opt_shape_3]
 *              |                                          |  |                                          |
 *              ---------------- profile 0 -----------------  ---------------- profile 1 -----------------
 * }
 *
 * The data after deserialization will be:
 * {
 *   tensor_a: {
 *     dim_0: [[min_shape_0, max_shape_0, opt_shape_0], [min_shape_1, max_shape_1, opt_shape_1]]
 *   },
 *   tensor_b: {
 *     dim_1: [[min_shape_2, max_shape_2, opt_shape_2], [min_shape_3, max_shape_3, opt_shape_3]]
 *   }
 * }
 */
std::unordered_map<std::string, std::unordered_map<size_t, std::vector<std::vector<int64_t>>>> DeserializeProfileV2(std::ifstream& infile) {
  ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] In DeserializeProfileV2()";
  // Load flexbuffer
  infile.seekg(0, std::ios::end);
  size_t length = infile.tellg();
  infile.seekg(0, std::ios::beg);
  std::unique_ptr<char[]> data{new char[length]};
  infile.read((char*)data.get(), length);
  infile.close();

  // Deserialize profile
  std::unordered_map<std::string, std::unordered_map<size_t, std::vector<std::vector<int64_t>>>> shape_ranges;
  auto tensors_range_entries = flexbuffers::GetRoot((const uint8_t*)data.get(), length).AsMap();
  auto keys = tensors_range_entries.Keys();
  auto values = tensors_range_entries.Values();
  for (size_t i = 0, end = keys.size(); i < end; ++i) {  // iterate tensors
    ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] input tensor is '" << keys[i].AsString().c_str() << "'";
    auto dim_range_vector = values[i].AsTypedVector();
    std::unordered_map<size_t, std::vector<std::vector<int64_t>>> inner_map;
    std::vector<std::vector<int64_t>> profile_vector;

    for (size_t k = 0; k < (dim_range_vector.size() / 4); k++) {  // iterate dim, min, max, opt for all profiles
      std::vector<int64_t> shape_vector;
      auto idx = 4 * k;
      auto dim = dim_range_vector[idx].AsInt64();
      shape_vector.push_back(dim_range_vector[idx + 1].AsInt64());  // min shape
      shape_vector.push_back(dim_range_vector[idx + 2].AsInt64());  // max shape
      shape_vector.push_back(dim_range_vector[idx + 3].AsInt64());  // opt shape

      if (inner_map.find(dim) == inner_map.end()) {
        inner_map[dim] = profile_vector;
      }
      inner_map[dim].push_back(shape_vector);
      ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] " << dim << ", " << shape_vector[0] << ", " << shape_vector[1] << ", " << shape_vector[2];
    }
    shape_ranges[keys[i].AsString().c_str()] = inner_map;
  }
  return shape_ranges;
}

/*
 * Compare profile shapes from profile file (.profile) with explicit profile min/max/opt shapes.
 * Return false meaning no need to rebuild engine if everything is same.
 * Otherwise return true and engine needs to be rebuilt.
 */
bool CompareProfiles(const std::string& file_name,
                     std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_min_shapes,
                     std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_max_shapes,
                     std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_opt_shapes) {
  std::ifstream profile_file(file_name, std::ios::binary | std::ios::in);
  if (!profile_file) {
    ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] " << file_name << " doesn't exist.";
    return true;
  }

  std::unordered_map<std::string, std::unordered_map<size_t, std::vector<std::vector<int64_t>>>> shape_ranges;
  shape_ranges = DeserializeProfileV2(profile_file);

  /* The format of the two data structures are below, for example:
   *
   * shape_ranges:
   * {
   *   tensor_a: {
   *     dim_0: [[min_shape, max_shape, opt_shape]],
   *     dim_2: [[min_shape, max_shape, opt_shape]]
   *   },
   *   tensor_b: {
   *     dim_1: [[min_shape, max_shape, opt_shape]]
   *   }
   * }
   *
   * profile_min_shapes:
   * {
   *   tensor_a: [[dim_0_value_0, dim_1_value_1, dim_2_value_2]],
   *   tensor_b: [[dim_0_value_3, dim_1_value_4, dim_2_value_5]]
   * }
   *
   */

  // Check number of dynamic shape inputs
  if (profile_min_shapes.size() != shape_ranges.size()) {
    ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Numbers of dynamic shape inputs are not the same.";
    return true;
  }

  // Iterate through shape_ranges map
  for (auto tensor_it = shape_ranges.begin(); tensor_it != shape_ranges.end(); tensor_it++) {  // iterate tensors
    auto tensor_name = tensor_it->first;
    if (profile_min_shapes.find(tensor_name) == profile_min_shapes.end()) {
      ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Tensor name '" << tensor_name << "' doesn't exist in trt_profile_min_shapes.";
      return true;
    }

    for (auto dim_it = tensor_it->second.begin(); dim_it != tensor_it->second.end(); dim_it++) {  // iterate dimensions
      auto dim = dim_it->first;
      auto num_profiles = GetNumProfiles(profile_min_shapes);

      if (dim_it->second.size() != static_cast<size_t>(num_profiles)) {
        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Numbers of profiles are not the same.";
        return true;
      }

      for (size_t i = 0; i < dim_it->second.size(); i++) {  // iterate (multiple) profile(s)
        auto shape_values = dim_it->second[i];
        if (dim > (profile_min_shapes[tensor_name][i].size() - 1)) {
          ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] dimension " << dim << " of '" << tensor_name << "' in " << file_name << " exceeds the total dimension of trt_profile_min_shapes.";
          return true;
        }

        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] min shape value of dimension " << dim << " of '" << tensor_name << "' is " << profile_min_shapes[tensor_name][i][dim];
        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] min shape value of dimension " << dim << " of '" << tensor_name << "' is " << shape_values[0] << " in " << file_name;
        if (profile_min_shapes[tensor_name][i][dim] != shape_values[0]) {
          ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] min shape values of dimension " << dim << " of '" << tensor_name << "' are not the same";
          return true;
        }

        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] max shape value of dimension " << dim << " of '" << tensor_name << "' is " << profile_max_shapes[tensor_name][i][dim];
        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] max shape value of dimension " << dim << " of '" << tensor_name << "' is " << shape_values[1] << " in " << file_name;
        if (profile_max_shapes[tensor_name][i][dim] != shape_values[1]) {
          ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] max shape values of dimension " << dim << " of '" << tensor_name << "' are not the same";
          return true;
        }

        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] opt shape value of dimension " << dim << " of '" << tensor_name << "' is " << profile_opt_shapes[tensor_name][i][dim];
        ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] opt shape value of dimension " << dim << " of '" << tensor_name << "' is " << shape_values[2] << " in " << file_name;
        if (profile_opt_shapes[tensor_name][i][dim] != shape_values[2]) {
          ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] opt shape values of dimension " << dim << " of '" << tensor_name << "' are not the same";
          return true;
        }
      }
    }
  }
  return false;
}

/*
 * Get cache by name
 *
 */
std::string GetCachePath(const std::string& root, const std::string& name) {
  if (root.empty()) {
    return name;
  } else {
    fs::path path = root;
    path.append(name);
    return path.string();
  }
}

/*
 * Get compute capability
 *
 */
std::string GetComputeCapacity(const cudaDeviceProp& prop) {
  const std::string compute_capability = std::to_string(prop.major * 10 + prop.minor);
  return compute_capability;
}

/*
 * Get Timing by compute capability
 *
 */
std::string GetTimingCachePath(const std::string& root, std::string& compute_cap) {
  // append compute capability of the GPU as this invalidates the cache and TRT will throw when loading the cache
  const std::string timing_cache_name = "TensorrtExecutionProvider_cache_sm" +
                                        compute_cap + ".timing";
  return GetCachePath(root, timing_cache_name);
}

/*
 * Get cache by type
 *
 * \param root root path of the cache
 * \param file_extension It could be ".engine", ".profile" or ".timing"
 */
std::vector<fs::path> GetCachesByType(const std::string& root, std::string file_extension) {
  std::vector<fs::path> cache_files;
  for (const auto& entry : fs::directory_iterator(root)) {
    if (fs::path(file_extension) == fs::path(entry).extension()) {
      cache_files.push_back(fs::path(entry));
    }
  }
  return cache_files;
}

bool IsCacheExistedByType(const std::string& root, std::string file_extension) {
  auto cache_files = GetCachesByType(root, file_extension);
  if (cache_files.size() == 0) {
    return false;
  }
  return true;
}

void RemoveCachesByType(const std::string& root, std::string file_extension) {
  auto cache_files = GetCachesByType(root, file_extension);
  for (const auto& entry : cache_files) {
    fs::remove(entry);
  }
}

bool ValidateProfileShapes(std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_min_shapes,
                           std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_max_shapes,
                           std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_opt_shapes) {
  if (profile_min_shapes.empty() && profile_max_shapes.empty() && profile_opt_shapes.empty()) {
    return true;
  }

  if ((profile_min_shapes.size() != profile_max_shapes.size()) &&
      (profile_min_shapes.size() != profile_opt_shapes.size()) &&
      (profile_max_shapes.size() != profile_opt_shapes.size())) {
    return false;
  }

  std::unordered_map<std::string, std::vector<std::vector<int64_t>>>::iterator it;
  for (it = profile_min_shapes.begin(); it != profile_min_shapes.end(); it++) {
    auto input_name = it->first;
    auto num_profile = it->second.size();

    // input_name must also be in max/opt profile
    if ((profile_max_shapes.find(input_name) == profile_max_shapes.end()) ||
        (profile_opt_shapes.find(input_name) == profile_opt_shapes.end())) {
      return false;
    }

    // number of profiles should be the same
    if ((num_profile != profile_max_shapes[input_name].size()) ||
        (num_profile != profile_opt_shapes[input_name].size())) {
      return false;
    }
  }

  return true;
}

/*
 * Make input-name and shape as a pair.
 * This helper function is being used by ParseProfileShapes().
 *
 * For example:
 * The input string is "input_id:32x1",
 * after the string is being parsed, the pair object is returned as below.
 * pair("input_id", [32, 1])
 *
 * Return true if string can be successfully parsed or false if string has wrong format.
 */
bool MakeInputNameShapePair(std::string pair_string, std::pair<std::string, std::vector<int64_t>>& pair) {
  if (pair_string.empty()) {
    return true;
  }

  ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] " << pair_string;

  std::stringstream input_string_stream(pair_string);
  char first_delim = ':';
  char second_delim = 'x';
  std::string input_name;
  std::string shape;
  std::getline(input_string_stream, input_name, first_delim);
  std::getline(input_string_stream, shape, first_delim);

  std::vector<int64_t> shapes;
  std::stringstream shape_string_stream(shape);
  std::string value;
  while (std::getline(shape_string_stream, value, second_delim)) {
    shapes.push_back(std::stoi(value));
  }

  // wrong input string
  if (input_name.empty() || shapes.empty()) {
    return false;
  }

  pair.first = input_name;
  pair.second = shapes;

  return true;
}

/*
 * Parse explicit profile min/max/opt shapes from TensorRT EP provider options.
 *
 * For example:
 * The provider option is --trt_profile_min_shapes="input_id:32x1,attention_mask:32x1,input_id:32x41,attention_mask:32x41",
 * after string is being parsed, the profile shapes has two profiles and is being represented as below.
 * {"input_id": [[32, 1], [32, 41]], "attention_mask": [[32, 1], [32, 41]]}
 *
 * Return true if string can be successfully parsed or false if string has wrong format.
 */
bool ParseProfileShapes(std::string profile_shapes_string, std::unordered_map<std::string, std::vector<std::vector<int64_t>>>& profile_shapes) {
  if (profile_shapes_string.empty()) {
    return true;
  }

  std::stringstream input_string_stream(profile_shapes_string);
  char delim = ',';
  std::string input_name_with_shape;  // input_name:shape, ex: "input_id:32x1"
  while (std::getline(input_string_stream, input_name_with_shape, delim)) {
    std::pair<std::string, std::vector<int64_t>> pair;
    if (!MakeInputNameShapePair(input_name_with_shape, pair)) {
      return false;
    }

    std::string input_name = pair.first;
    if (profile_shapes.find(input_name) == profile_shapes.end()) {
      std::vector<std::vector<int64_t>> profile_shape_vector;
      profile_shapes[input_name] = profile_shape_vector;
    }
    profile_shapes[input_name].push_back(pair.second);

    ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] " << input_name;
    std::string shape_string = "";
    for (auto v : pair.second) {
      shape_string += std::to_string(v);
      shape_string += ", ";
    }
    ////LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] " << shape_string;
  }

  return true;
}

std::vector<std::string> split(const std::string& str, char delimiter) {
  std::vector<std::string> tokens;
  std::string token;
  std::istringstream tokenStream(str);
  while (std::getline(tokenStream, token, delimiter)) {
    tokens.push_back(token);
  }
  return tokens;
}

std::string join(const std::vector<std::string>& vec, const std::string& delimiter) {
  std::string result;
  for (size_t i = 0; i < vec.size(); ++i) {
    result += vec[i];
    if (i < vec.size() - 1) {
      result += delimiter;
    }
  }
  return result;
}

/*
 * Parse engine cache name suffix when user customizes prefix for engine cache name
 *
 * For example:
 * When default subgraph name is "TensorrtExecutionProvider_TRTKernel_graph_torch-jit-export_2068723788287043730_189_189_fp16"
 * This func will generate the suffix "2068723788287043730_189_fp16"
 *
 */
std::string GetCacheSuffix(const std::string& fused_node_name, const std::string& trt_node_name_with_precision) {
  std::vector<std::string> split_fused_node_name = split(fused_node_name, '_');
  if (split_fused_node_name.size() >= 3) {
    // Get index of model hash from fused_node_name
    std::string model_hash = split_fused_node_name[split_fused_node_name.size() - 3];
    size_t index = fused_node_name.find(model_hash);
    // Parse suffix from trt_node_name_with_precision, as it has additional precision info
    std::vector<std::string> suffix_group = split(trt_node_name_with_precision.substr(index), '_');
    if (suffix_group.size() > 2) {
      suffix_group.erase(suffix_group.begin() + 2);
    }
    return join(suffix_group, "_");
  }
  return "";
}

}  // namespace trt_ep
